package com.tianhai.db;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * 锁管理器 - 实现两阶段锁协议（2PL）
 * 支持共享锁（SHARED）和排他锁（EXCLUSIVE）
 * 包含死锁预防机制（超时）
 */
public class LockManager {
    
    // 锁类型
    public enum LockType {
        SHARED,    // 共享锁（读锁）
        EXCLUSIVE  // 排他锁（写锁）
    }
    
    // 资源锁表：<资源ID, 锁持有信息>
    private final Map<String, ResourceLock> lockTable;
    
    // 事务锁持有记录：<事务ID, 持有的资源集合>
    private final Map<Long, Set<String>> transactionLocks;
    
    // 全局读写锁
    private final ReentrantReadWriteLock globalLock;
    
    // 死锁预防：锁等待超时时间（毫秒）
    private final long lockTimeout;
    
    /**
     * 构造函数
     * @param lockTimeout 锁等待超时时间（毫秒），0表示无限等待
     */
    public LockManager(long lockTimeout) {
        this.lockTable = new ConcurrentHashMap<>();
        this.transactionLocks = new ConcurrentHashMap<>();
        this.globalLock = new ReentrantReadWriteLock();
        this.lockTimeout = lockTimeout;
    }
    
    /**
     * 默认构造函数（5秒超时）
     */
    public LockManager() {
        this(5000);
    }
    
    /**
     * 获取锁
     * @param transactionId 事务ID
     * @param resource 资源标识
     * @param lockType 锁类型
     * @return true表示成功获取锁，false表示超时
     */
    public boolean acquireLock(long transactionId, String resource, LockType lockType) 
            throws InterruptedException {
        
        globalLock.writeLock().lock();
        try {
            // 获取或创建资源锁
            ResourceLock resourceLock = lockTable.computeIfAbsent(
                resource, k -> new ResourceLock(resource)
            );
            
            // 检查是否已持有锁
            if (resourceLock.isHeldBy(transactionId)) {
                // 已持有锁，检查是否需要升级
                if (resourceLock.getLockType(transactionId) == LockType.SHARED && 
                    lockType == LockType.EXCLUSIVE) {
                    // 锁升级：共享锁 -> 排他锁
                    return upgradeLock(transactionId, resource, resourceLock);
                }
                // 已持有相同或更强的锁
                return true;
            }
            
            // 尝试获取新锁
            long startTime = System.currentTimeMillis();
            long deadline = lockTimeout > 0 ? startTime + lockTimeout : Long.MAX_VALUE;
            
            while (!canGrantLock(resourceLock, lockType)) {
                // 检查超时
                if (System.currentTimeMillis() >= deadline) {
                    return false; // 超时，死锁预防
                }
                
                // 等待一小段时间后重试
                globalLock.writeLock().unlock();
                Thread.sleep(10);
                globalLock.writeLock().lock();
                
                // 重新获取资源锁（可能已改变）
                resourceLock = lockTable.get(resource);
                if (resourceLock == null) {
                    resourceLock = lockTable.computeIfAbsent(
                        resource, k -> new ResourceLock(resource)
                    );
                }
            }
            
            // 授予锁
            resourceLock.grant(transactionId, lockType);
            
            // 记录事务持有的锁
            transactionLocks.computeIfAbsent(transactionId, k -> new HashSet<>())
                           .add(resource);
            
            return true;
            
        } finally {
            globalLock.writeLock().unlock();
        }
    }
    
    /**
     * 释放锁
     * @param transactionId 事务ID
     * @param resource 资源标识
     */
    public void releaseLock(long transactionId, String resource) {
        globalLock.writeLock().lock();
        try {
            ResourceLock resourceLock = lockTable.get(resource);
            if (resourceLock != null) {
                resourceLock.release(transactionId);
                
                // 如果资源上没有锁了，移除资源锁对象
                if (resourceLock.isEmpty()) {
                    lockTable.remove(resource);
                }
            }
            
            // 更新事务锁记录
            Set<String> locks = transactionLocks.get(transactionId);
            if (locks != null) {
                locks.remove(resource);
                if (locks.isEmpty()) {
                    transactionLocks.remove(transactionId);
                }
            }
            
        } finally {
            globalLock.writeLock().unlock();
        }
    }
    
    /**
     * 释放事务的所有锁
     * @param transactionId 事务ID
     */
    public void releaseAllLocks(long transactionId) {
        globalLock.writeLock().lock();
        try {
            Set<String> locks = transactionLocks.get(transactionId);
            if (locks != null) {
                // 创建副本避免并发修改
                Set<String> locksCopy = new HashSet<>(locks);
                for (String resource : locksCopy) {
                    releaseLock(transactionId, resource);
                }
            }
        } finally {
            globalLock.writeLock().unlock();
        }
    }
    
    /**
     * 检查是否可以授予锁（锁兼容性检查）
     */
    private boolean canGrantLock(ResourceLock resourceLock, LockType requestType) {
        if (resourceLock.isEmpty()) {
            return true; // 无锁持有
        }
        
        // 锁兼容性矩阵：
        // SHARED 与 SHARED 兼容
        // SHARED 与 EXCLUSIVE 不兼容
        // EXCLUSIVE 与任何锁都不兼容
        
        if (requestType == LockType.SHARED) {
            // 请求共享锁：只要没有排他锁即可
            return !resourceLock.hasExclusiveLock();
        } else {
            // 请求排他锁：必须无任何锁
            return false;
        }
    }
    
    /**
     * 锁升级：共享锁 -> 排他锁
     */
    private boolean upgradeLock(long transactionId, String resource, ResourceLock resourceLock) 
            throws InterruptedException {
        
        long startTime = System.currentTimeMillis();
        long deadline = lockTimeout > 0 ? startTime + lockTimeout : Long.MAX_VALUE;
        
        // 等待其他共享锁释放
        while (resourceLock.getSharedCount() > 1) {
            if (System.currentTimeMillis() >= deadline) {
                return false; // 超时
            }
            
            globalLock.writeLock().unlock();
            Thread.sleep(10);
            globalLock.writeLock().lock();
        }
        
        // 升级为排他锁
        resourceLock.upgrade(transactionId);
        return true;
    }
    
    /**
     * 获取事务持有的锁数量
     */
    public int getLockCount(long transactionId) {
        globalLock.readLock().lock();
        try {
            Set<String> locks = transactionLocks.get(transactionId);
            return locks != null ? locks.size() : 0;
        } finally {
            globalLock.readLock().unlock();
        }
    }
    
    /**
     * 检查事务是否持有资源的锁
     */
    public boolean holdsLock(long transactionId, String resource) {
        globalLock.readLock().lock();
        try {
            ResourceLock resourceLock = lockTable.get(resource);
            return resourceLock != null && resourceLock.isHeldBy(transactionId);
        } finally {
            globalLock.readLock().unlock();
        }
    }
    
    /**
     * 资源锁
     */
    private static class ResourceLock {
        private final String resource;
        private final Map<Long, LockType> holders; // 锁持有者：<事务ID, 锁类型>
        
        public ResourceLock(String resource) {
            this.resource = resource;
            this.holders = new HashMap<>();
        }
        
        /**
         * 授予锁
         */
        public void grant(long transactionId, LockType lockType) {
            holders.put(transactionId, lockType);
        }
        
        /**
         * 释放锁
         */
        public void release(long transactionId) {
            holders.remove(transactionId);
        }
        
        /**
         * 升级锁
         */
        public void upgrade(long transactionId) {
            holders.put(transactionId, LockType.EXCLUSIVE);
        }
        
        /**
         * 检查事务是否持有锁
         */
        public boolean isHeldBy(long transactionId) {
            return holders.containsKey(transactionId);
        }
        
        /**
         * 获取事务持有的锁类型
         */
        public LockType getLockType(long transactionId) {
            return holders.get(transactionId);
        }
        
        /**
         * 是否为空（无锁持有）
         */
        public boolean isEmpty() {
            return holders.isEmpty();
        }
        
        /**
         * 是否有排他锁
         */
        public boolean hasExclusiveLock() {
            for (LockType type : holders.values()) {
                if (type == LockType.EXCLUSIVE) {
                    return true;
                }
            }
            return false;
        }
        
        /**
         * 获取共享锁数量
         */
        public int getSharedCount() {
            int count = 0;
            for (LockType type : holders.values()) {
                if (type == LockType.SHARED) {
                    count++;
                }
            }
            return count;
        }
        
        @Override
        public String toString() {
            return "ResourceLock{resource=" + resource + ", holders=" + holders + "}";
        }
    }
}
